package org.paraj.mapreduce.executors;

import com.google.inject.Inject;
import org.paraj.mapreduce.ResultAggregator;

import java.io.Serializable;
import java.util.*;

public class InMemoryResultService implements ResultService{
    private Map<UUID, TaskCoordinator> tasks = new HashMap<UUID, TaskCoordinator>();
    private Distributor distributor;


    public void startingTask(UUID parentId, UUID uid) {
        System.out.println("Starting " + uid);
        tasks.put(uid, new TaskCoordinator( uid, parentId));
    }

    public void initializeComplexTask(UUID uid, int subtaskCount) {
        tasks.get(uid).setSubtasksCount(subtaskCount);
    }

    public void finishedSubtask(UUID parentTaskId, Object result) {
        tasks.get(parentTaskId).addSubtaskResult(result);
    }

    public <R> R waitForResult(UUID uid) throws InterruptedException {
        return (R)tasks.get(uid).waitForResult();
    }

    public void registerAggregator(UUID uid, ResultAggregator resultAggregator) {
        tasks.get(uid).setAggregator(resultAggregator);
    }

    public Collection waitForSubResults(UUID uid) throws InterruptedException {
        return tasks.get(uid).waitForSubResults();
    }

    public synchronized Collection getSubtaskResults(UUID taskId){
        return tasks.get(taskId).subtaskResults;
    }

    public void finishedTask(UUID taskId, Object aggregatedResult) {
        System.out.println("Finished " + taskId+", result="+aggregatedResult);
        tasks.get(taskId).finished(aggregatedResult);
    }



    private void schedule(Runnable runnable) {
        distributor.schedule(runnable);
    }

    public void setDistributor(Distributor distributor) {
        this.distributor = distributor;
    }

    private class TaskCoordinator {
        int subtaskCount;
        private UUID uid;
        private UUID parentId;
        private Collection subtaskResults;
        private Object aggregatedResult;
        private boolean finished;
        private boolean subtasksFinished;
        private ResultAggregator resultAggregator;

        public TaskCoordinator(UUID uid, UUID parentId) {
            this.uid = uid;
            this.parentId = parentId;
            finished = false;
            subtaskResults = new ArrayList();
        }

        public  void addSubtaskResult(Object result) {
            if (addResultAndGetSize(result) == subtaskCount) onAllSubtasksFinished();
        }

        private synchronized int addResultAndGetSize(Object result) {
            subtaskResults.add(result);
            return subtaskResults.size();
        }

        private synchronized void onAllSubtasksFinished() {
            if (resultAggregator != null) schedule(new RunAggregation(uid, resultAggregator));
            subtasksFinished = true;
            notifyAll();
        }

        public synchronized <R> R waitForResult() throws InterruptedException {
            while(!finished) wait();
            return (R)aggregatedResult;
        }

        public synchronized void finished(Object aggregatedResult) {
            if (parentId != null) {
                finishedSubtask(parentId, aggregatedResult);
                subtaskResults = null;
                this.aggregatedResult = null;
            }
            this.aggregatedResult = aggregatedResult;
            finished = true;
            notifyAll();
        }

        public void setAggregator(ResultAggregator resultAggregator) {
            this.resultAggregator = resultAggregator;
        }

        public Collection waitForSubResults() throws InterruptedException {
            while(!subtasksFinished) wait();
            return subtaskResults;
        }

        public void setSubtasksCount(int subtaskCount) {
            this.subtaskCount = subtaskCount;
        }
    }

    private static class RunAggregation implements Runnable, Serializable{
        private final ResultAggregator resultAggregator;
        private UUID taskId;
        @Inject transient ResultService resultService;

        public RunAggregation(UUID taskId, ResultAggregator resultAggregator){
            this.resultAggregator = resultAggregator;
            this.taskId = taskId;
        }

        public void setResultService(ResultService resultService) {
            this.resultService = resultService;
        }

        public void run(){
            Object aggregatedResult = resultAggregator.aggregate(resultService.getSubtaskResults(taskId));
            resultService.finishedTask(taskId, aggregatedResult);
        }
    }


}
